import numpy as np
from copy import deepcopy
from math import ceil, modf
from scipy.integrate import odeint

dWdN = 1e6 * 24 * 3600 / (1.23 / 235)

def Power2phi(power, N_U235, N_U239, N_Pu239, N_Pu241, N_Am241):
    return power / ((585 * N_U235 + 15 * N_U239 + 750 * N_Pu239 + 1010 * N_Pu241 + 3 * N_Am241) * dWdN)

class Calculate:
    
    def __init__(self, data, setting):
        self.CHAIN_NAMES = ['235U', '238U']
        self.FUEL = {}
        for name in self.CHAIN_NAMES:
            self.FUEL[name] = {}
            for nuclide, weight_ratio in setting['Initial'][name].items():
                self.FUEL[name][nuclide] = weight_ratio / (int(nuclide[0:3]) * 1e-3)
        self.U235 = self.FUEL['235U']['235U']
        self.DATA = data
        self.DT = setting['DT'] * 24 * 60 * 60
        self.METHOD = setting['Method']
        if self.METHOD == 'Flux':
            self.phi = setting['Flux'] * 1e-24
        else:
            self.POWER = setting['Power'] * 1e6

        self.numerical_result = None
        
    def Find_Nuclide(self, Data, nuclide_id):
        nuclides = Data[:, 0]
        return [index for index, iterm in enumerate(nuclides) if int(iterm) == int(nuclide_id)][0]

    def Cal_A(self, row):
        if row[4] == 'Neutron Absorbtion':
            return -self.phi*(float(row[5]) + float(row[6]))
        elif row[4] == 'Neutron Absorbtion and Beta Decay':
            return -self.phi*(float(row[5]) + float(row[6])) - float(row[7])

    def Cal_A_(self, parent_row, tf_type):
        if tf_type == 'Neutron Absorbtion':
            return self.phi*float(parent_row[6])
        elif tf_type == 'Beta Decay':
            return float(parent_row[7])

class Analytic(Calculate):
    
    def __init__(self, data, setting):
        super().__init__(data, setting)
        self.analytic_result = None

    def Stat_Gen(self, Data, row_id, gen=1):
        gen += 1
        if int(Data[row_id][2]) == 0:
            return gen
        else:
            parent_nuclide = Data[row_id][2]
            return self.Stat_Gen(Data, self.Find_Nuclide(Data, parent_nuclide), gen)

    def Step(self):
        result_A = {
            '235U': [],
            '238U': []
        }
        result_C = {
            '235U': [],
            '238U': []
        }
        result_name = {
            '235U': [],
            '238U': []
        }
        result_Fuel = {
            '235U': {},
            '238U': {}
        }

        for name in self.CHAIN_NAMES:
            Data = self.DATA[name]
            # PDF上的解法有bug，当分支不是在末端时，必须回溯父代参数，而非39页式(2)求和即可。对于目前简单情况，可以使用。
            for row in Data:
                result_A[name].append(self.Cal_A(row))
                result_C[name].append([])
                result_name[name].append(row[1])
                if row[3] == 'Original':
                    gen_id = 1
                    result_C[name][-1].append(self.FUEL[name]
                                                [row[1]])
                else:
                    parent_id = self.Find_Nuclide(Data, row[2])
                    gen_id = self.Stat_Gen(Data, parent_id)
                    C_ii = self.FUEL[name][row[1]]
                    A_ = self.Cal_A_(Data[parent_id], row[3])
                    for C_id in range(gen_id - 1):
                        result_C[name][-1].append(A_ * result_C[name][-2][C_id] / (
                            result_A[name][C_id] - result_A[name][-1]))
                        C_ii -= result_C[name][-1][-1]
                    result_C[name][-1].append(C_ii)

                fuel = 0
                index = 0
                for A_jj, C_ij in zip(result_A[name], result_C[name][-1]):
                    index += 1
                    if index < gen_id:
                        fuel += C_ij * np.exp(A_jj * self.DT)
                    else:
                        fuel += C_ij * np.exp(result_A[name][-1] * self.DT)
                result_Fuel[name][row[1]] = fuel

        self.FUEL = result_Fuel
        return result_A, result_C, result_name

    def Cal_Numerical(self, analytic_result, index, TIME_SEQUENCE, t):
        time_sequence = deepcopy(TIME_SEQUENCE)
        for name in self.CHAIN_NAMES:
            for index_0, C_i in enumerate(analytic_result['C'][name]):
                N_sequence = np.zeros(len(time_sequence))
                N_t = 0
                for index_1, C_ij in enumerate(C_i):
                    if index_1 == len(C_i) - 1:
                        N_sequence += C_ij * \
                            np.exp(analytic_result['A'][name]
                                   [index_0] * time_sequence)
                        N_t += C_ij * \
                            np.exp(analytic_result['A'][name][index_0] * t)
                    else:
                        N_sequence += C_ij * \
                            np.exp(analytic_result['A'][name]
                                   [index_1] * time_sequence)
                        N_t += C_ij * \
                            np.exp(analytic_result['A'][name][index_1] * t)

                nuclide_name = analytic_result['Name'][name][index_0]
                if index == 0:
                    self.numerical_result['Data'][name][nuclide_name] = N_sequence.tolist(
                    )
                else:
                    self.numerical_result['Data'][name][nuclide_name] = np.hstack(
                        (np.array(self.numerical_result['Data'][name][nuclide_name]), N_sequence)).tolist()
                self.numerical_result['t'][1][name][nuclide_name] = N_t

        if index == 0:
            self.numerical_result['Data']['t'] = time_sequence.tolist()
        else:
            time_sequence += index * self.DT
            self.numerical_result['Data']['t'] = np.hstack(
                (np.array(self.numerical_result['Data']['t']), time_sequence)).tolist()

    def main(self, t, precision=1e-3):
        self.analytic_result = {
            't': t,
            'Data': []
        }
        self.numerical_result = {
            'Method': self.METHOD,
            't': [
                t,
                {
                    '235U': {},
                    '238U': {}
                }
            ],
            'Data': {
                '235U': {},
                '238U': {}
            }
        }
        if self.METHOD == 'Power':
            self.numerical_result['Power'] = self.POWER * 1e-6

        steps = int(t / self.DT) + 1
        t = modf(t / self.DT)[0]
        time_sequence = np.linspace(
            0, self.DT, num=ceil(1/precision), endpoint=False)
        for index in range(steps):
            step_result = {
                'Fuel': self.FUEL
            }

            if float(self.FUEL['235U']['235U']) < 0.1 * self.U235:
                break
            self.phi = Power2phi(
                self.POWER, self.FUEL['235U']['235U'], self.FUEL['238U']['239U'], self.FUEL['238U']['239Pu'], self.FUEL['238U']['241Pu'], self.FUEL['238U']['241Am']) if self.METHOD == 'Power' else self.phi

            step_result['A'], step_result['C'], step_result['Name'] = self.Step()
            self.analytic_result['Data'].append(step_result)
            self.Cal_Numerical(step_result, index, time_sequence, t)

class Numerical(Calculate):

    def __init__(self, data, setting):
        super().__init__(data, setting)
    
    def Step(self, index, TIME_SEQUENCE):
        time_sequence = deepcopy(TIME_SEQUENCE)
        result_A = {
            '235U': [],
            '238U': []
        }
        result_C = {
            '235U': [],
            '238U': []
        }
        result_name = {
            '235U': [],
            '238U': []
        }
        result_A_ = {
            '235U': [],
            '238U': []
        }
        
        for name in self.CHAIN_NAMES:
            Data = self.DATA[name]
            # PDF上的解法有bug，当分支不是在末端时，必须回溯父代参数，而非39页式(2)求和即可。对于目前简单情况，可以使用。
            for row in Data:
                result_A[name].append(self.Cal_A(row))
                result_C[name].append(self.FUEL[name][row[1]])
                result_name[name].append(row[1])
                if row[3] == 'Original':
                    result_A_[name].append(0)
                else:
                    parent_id = self.Find_Nuclide(Data, row[2])
                    result_A_[name].append(self.Cal_A_(Data[parent_id], row[3]))

            def Cal_Numerical(N, t):
                dNdt = []
                for n_id, n in enumerate(N):
                    if n_id == 0:
                        dNdt.append(result_A[name][n_id] * n)
                    else:
                        dNdt.append(result_A_[name][n_id] * N[n_id - 1] + result_A[name][n_id] * n)
                return dNdt
            
            sol = odeint(Cal_Numerical, result_C[name], time_sequence)
            for name_id, nuclide_name in enumerate(result_name[name]):
                if index == 0:
                    self.numerical_result['Data'][name][nuclide_name] = sol[:, name_id].tolist()
                else:
                    self.numerical_result['Data'][name][nuclide_name] = np.hstack((np.array(self.numerical_result['Data'][name][nuclide_name]), sol[:, name_id])).tolist()
                self.FUEL[name][nuclide_name] = sol[:, name_id][-1]

        if index == 0:
            self.numerical_result['Data']['t'] = time_sequence.tolist()
        else:
            time_sequence += index * self.DT
            self.numerical_result['Data']['t'] = np.hstack(
                (np.array(self.numerical_result['Data']['t']), time_sequence)).tolist()

    def main(self, t, precision=1e-3):
        self.numerical_result = {
            'Method': self.METHOD,
            't': [
                t,
                {
                    '235U': {},
                    '238U': {}
                }
            ],
            'Data': {
                '235U': {},
                '238U': {}
            }
        }
        if self.METHOD == 'Power':
            self.numerical_result['Power'] = self.POWER * 1e-6
            
        steps = int(t / self.DT) + 1
        t = modf(t / self.DT)[0]
        time_sequence = np.linspace(
            0, self.DT, num=ceil(1/precision), endpoint=True)
        for index in range(steps):

            if float(self.FUEL['235U']['235U']) < 0.1 * self.U235:
                break
            self.phi = Power2phi(
                self.POWER, self.FUEL['235U']['235U'], self.FUEL['238U']['239U'], self.FUEL['238U']['239Pu'], self.FUEL['238U']['241Pu'], self.FUEL['238U']['241Am']) if self.METHOD == 'Power' else self.phi
            self.Step(index, time_sequence)